import torch

# ---------------------- 1. 模拟初始化缓存（对应k_cache） ----------------------
# 假设：缓存分2块（num_blocks=2），每块存3个token（block_size=3），1个头（num_heads=1），特征维度=2
block_size = 3
k_cache = torch.zeros((2, block_size, 1, 2), dtype=torch.float32)
print("初始化后的k_cache（全零）：")
print(k_cache)
# 初始化形状：(2块, 每块3token, 1头, 2维特征) → 对应你的(2,64,32,64)简化版
# 输出：
# tensor([[[[0., 0.]], [[0., 0.]], [[0., 0.]]],  # 第0块：3个token，全零
#         [[[0., 0.]], [[0., 0.]], [[0., 0.]]]]) # 第1块：3个token，全零


# ---------------------- 2. 模拟当前计算的K特征（对应k） ----------------------
# 假设：当前处理4个token（seq_len=4），1个头，特征维度=2
k = torch.tensor([
    [[1.1, 1.2]],  # 第0个token的K特征
    [[2.1, 2.2]],  # 第1个token的K特征
    [[3.1, 3.2]],  # 第2个token的K特征
    [[4.1, 4.2]]   # 第3个token的K特征
], dtype=torch.float32)
print("\n当前要写入的k（4个token的K特征）：")
print(k)
# k形状：(4token, 1头, 2维特征) → 对应你的(100,32,64)简化版


# ---------------------- 3. 计算缓存索引（对应block_idx和token_idx_in_block） ----------------------
# 模拟slot_mapping：每个token在缓存中的“全局索引”（比如第0~3个token对应全局索引0~3）
slot_mapping = torch.tensor([0, 1, 2, 3])  # 4个token的全局索引
# 计算每个token属于“哪一块”（block_idx）和“块内哪个位置”（token_idx_in_block）
block_idx = slot_mapping // block_size  # 整除：0//3=0, 1//3=0, 2//3=0, 3//3=1 → 前3个在第0块，第4个在第1块
token_idx_in_block = slot_mapping % block_size  # 取余：0%3=0,1%3=1,2%3=2,3%3=0 → 块内位置

print("\n计算出的索引：")
print(f"block_idx（属于哪一块）：{block_idx}")  # 输出：tensor([0, 0, 0, 1])
print(f"token_idx_in_block（块内位置）：{token_idx_in_block}")  # 输出：tensor([0, 1, 2, 0])


# ---------------------- 4. 批量更新缓存（核心步骤！） ----------------------
# 用k的值，批量覆盖k_cache中“block_idx+token_idx_in_block”定位的位置
k_cache[block_idx, token_idx_in_block] = k

print("\n更新后的k_cache（指定位置被k覆盖）：")
print(k_cache)
# 输出结果分析：
# 第0块的0、1、2位置 → 被k[0]、k[1]、k[2]覆盖（值为1.1/1.2, 2.1/2.2, 3.1/3.2）
# 第1块的0位置 → 被k[3]覆盖（值为4.1/4.2）
# 第1块的1、2位置 → 仍为0（未被更新）